## code to run MSA-level MSMs (using median ZRI) and other models ##
## with non-imputed data ##

msm_noimp_med <- function(dv1, dv2, dvp, num){
  
  af.msa.done <- boot.msa.noimp.data
  
  prop.table(table(af.msa.done$treat_2006,useNA="ifany"))
  
  print(paste0("MSA ", num))
  
  vars.2022 <- paste(c("pop_total_2022","pop_density_2022",
                       "per_18t34_2022","per_35t64_2022",
                       "per_65over_2022","entropy_er_2022",
                       "per_child_2022", "per_nomove_2022",
                       "per_fb_2022","hownrt_2022",
                       "median_yr_str_2022","median_hhld_inc_2022",
                       "hhld_pov_rt_2022","unemr_2022",
                       "gini_2022"
  ), collapse = "+")
  
  
  vars.2009 <- paste(c("pop_total_2009","pop_density_2009",
                       "per_18t34_2009","per_35t64_2009",
                       "per_65over_2009","entropy_er_2009",
                       "per_child_2009", "per_nomove_2009",
                       "per_fb_2009","hownrt_2009",
                       "median_yr_str_2009","median_hhld_inc_2009",
                       "hhld_pov_rt_2009","unemr_2009",
                       "gini_2009"
  ), collapse = "+")
  
  fe.vars <- paste(c("pop_total","pop_density",
                     "per_18t34","per_35t64",
                     "per_65over","entropy_er",
                     "per_child", "per_nomove",
                     "per_fb","hownrt",
                     "median_yr_str","median_hhld_inc",
                     "hhld_pov_rt","unemr",
                     "gini","munis_tot",
                     "CBSA","tp"
  ), collapse = "+")
  
  ## regression formulas ##
  
  form1 <- paste0("treat_2022 ~ ",
                  dv1, " + treat_2006 +",
                  vars.2022, "+ munis_tot")
  
  form2 <- paste0("treat_2006 ~ ", 
                  vars.2009, "+ munis_tot")
  
  form3a <- paste0(dv2," ~",
                   "thist + munis_tot")
  
  form3b <- paste0(dv2," ~",
                   "thist")
  
  form4 <- paste0(dv2," ~",
                  "zri_2022_median_st +",
                  "zri_2006_median_st +",
                  dv1," +",
                  vars.2022," +",
                  vars.2009, 
                  "+ munis_tot")
  
  form5 <- paste0(dvp," ~",
                  "zri_median"," +",
                  fe.vars)
  
  dv1 <- ensym(dv1)
  
  msa.2022.comp <- af.msa.done %>%
    filter(!is.na(!!dv1))
  
  if(nrow(msa.2022.comp) < 30){
    
    print("TOO FEW OBS, RETURNING NA")
    return(NA)
  }
  
  msa.2022 <- glm(as.formula(form1),
                  data = msa.2022.comp,
                  binomial(link = 'logit'))
  
  summary(msa.2022)
  
  ## predicted probabilities ##
  
  ppout.2022 <- predict(msa.2022, type = "response", se.fit = TRUE)
  pps.2022 <- ppout.2022$fit
  
  ## numerator for stabilized weights ##
  
  mean(msa.2022.comp$treat_2022)
  pps.2022.num <- rep(mean(msa.2022.comp$treat_2022),
                      length(pps.2022))
  
  ## now 2006 ##
  
  msa.2006.comp <- af.msa.done %>%
    filter(!is.na(treat_2006) & 
             CBSA %in% msa.2022.comp$CBSA)
  
  msa.2006 <- glm(as.formula(form2),
                  data = msa.2006.comp,
                  binomial(link = 'logit'))
  
  summary(msa.2006)
  
  ppout.2006 <- predict(msa.2006, type = "response", se.fit = TRUE)
  pps.2006 <- ppout.2006$fit

  mean(msa.2006.comp$treat_2006,na.rm=T)
  pps.2006.num <- rep(mean(msa.2006.comp$treat_2006,na.rm=T),
                      length(pps.2006))
  
  ## combine the weights ##
  
  print("LENGTHS - NOIMP")
  print(length(pps.2006))
  print(length(pps.2022))
  
  ## attach the pps and numerators and create weights ##
  
  af.msa.wts <- msa.2022.comp %>%
    mutate(thist = treat_2006 + treat_2022,
           num_t_2006 = pps.2006.num,
           pps_2006 = pps.2006,
           num_t_2022 = pps.2022.num,
           pps_2022 = pps.2022,
           num_2006 = ifelse(treat_2006 == 1, num_t_2006, 1-num_t_2006),
           den_2006 = ifelse(treat_2006 == 1, pps_2006, 1-pps_2006),
           num_2022 = ifelse(treat_2022 == 1, num_t_2022, 1-num_t_2022),
           den_2022 = ifelse(treat_2022 == 1, pps_2022, 1-pps_2022),
           swt_2006 = num_2006/den_2006,
           swt_2022 = num_2022/den_2022,
           swt_fin = swt_2006 * swt_2022)
  
  ## truncate the weights ##
  
  wts.cb.p1 <- quantile(af.msa.wts$swt_fin, 0.01, na.rm=T)
  wts.cb.p99 <- quantile(af.msa.wts$swt_fin,0.99, na.rm=T)
  
  af.msa.wts$swt_fin[af.msa.wts$swt_fin > wts.cb.p99] <- wts.cb.p99
  af.msa.wts$swt_fin[af.msa.wts$swt_fin < wts.cb.p1] <- wts.cb.p1
  summary(af.msa.wts$swt_fin)
  sd(af.msa.wts$swt_fin, na.rm=T)
  
  ## assign the weights for export ##
  wts.fin <- af.msa.wts$swt_fin
  
  ## msm ##
  
  msa.msm <- lm(as.formula(form3a),
                data = af.msa.wts,
                weights=swt_fin)
  
  ## unadjusted ## 
  
  msa.noadj <- lm(as.formula(form3b),
                  data = af.msa.wts)
  
  ## adl ## 
  
  msa.adl <- lm(as.formula(form4),
                data = af.msa.wts)
  
  ## fixed effects ## 
  
  msa.fe <- lm(as.formula(form5),
               data = af.msa.fedatanoimp)
  
  out.msm <- summary(msa.msm)
  out.noadj <- summary(msa.noadj)
  out.adl <- summary(msa.adl)
  out.fe <- summary(msa.fe)
  
  
  ## assess balance ##
  
  af.msa.bal <- af.msa.wts %>%
    filter(!is.na(treat_2006))
  
  bal.t2.vars <- c("pop_total_2022","pop_density_2022",
                   "per_18t34_2022","per_35t64_2022",
                   "per_65over_2022","entropy_er_2022",
                   "per_child_2022", "per_nomove_2022",
                   "per_fb_2022","hownrt_2022",
                   "median_yr_str_2022","median_hhld_inc_2022",
                   "hhld_pov_rt_2022","unemr_2022",
                   "gini_2022")
  
  bal.t1.vars <- c("pop_total_2009","pop_density_2009",
                   "per_18t34_2009","per_35t64_2009",
                   "per_65over_2009","entropy_er_2009",
                   "per_child_2009", "per_nomove_2009",
                   "per_fb_2009","hownrt_2009",
                   "median_yr_str_2009","median_hhld_inc_2009",
                   "hhld_pov_rt_2009","unemr_2009",
                   "gini_2009")
  
  bal.t2 <- function(var){
    
    unwt2 <- lm(as.formula(paste0(var,
                                  " ~ treat_2022 + treat_2006 + munis_tot")),
                data = af.msa.bal)
    
    out2a <- summary(unwt2)
    
    
    wt2 <- lm(as.formula(paste0(var,
                                " ~ treat_2022 + treat_2006 + munis_tot")),
              data = af.msa.bal,
              weights=swt_fin)
    
    out2b <- summary(wt2)
    
    outres <- list(out2a,out2b)
    
    return(outres)
    
  }
  
  t2.bal <- lapply(bal.t2.vars, bal.t2)  
  
  bal.t1 <- function(var){
    
    unwt1 <- lm(as.formula(paste0(var,
                                  " ~ treat_2006 + munis_tot")),
                data = af.msa.bal)
    
    out1a <- summary(unwt1)
    
    
    wt1 <- lm(as.formula(paste0(var,
                                " ~ treat_2006 + munis_tot")),
              data = af.msa.bal,
              weights=swt_fin)
    
    out1b <- summary(wt1)
    
    outres <- list(out1a,out1b)
    
    return(outres)
    
  }
  
  t1.bal <- lapply(bal.t1.vars, bal.t1)  
  
  ## capture all output ## 
  
  res <- list(out.msm,
              wts.fin,
              out.noadj,
              out.adl,
              out.fe,
              t2.bal,
              t1.bal)
  
  return(res)
  
}

